{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp test_base\n",
    "import os\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Base\n",
    "\n",
    "Basic setup for testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "import shutil\n",
    "import tempfile\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from loguru import logger\n",
    "from m3tl.input_fn import train_eval_input_fn\n",
    "from m3tl.model_fn import BertMultiTask, BertMultiTaskBody, BertMultiTaskTop\n",
    "from m3tl.params import Params\n",
    "from m3tl.predefined_problems import (\n",
    "    get_fake_contrastive_learning_fn,\n",
    "    get_weibo_fake_cls_fn,\n",
    "    get_weibo_fake_multi_cls_fn,\n",
    "    get_weibo_fake_ner_fn,\n",
    "    get_weibo_fake_regression,\n",
    "    get_weibo_fake_vector_fit,\n",
    "    get_weibo_masklm, get_weibo_premask_mlm,\n",
    "    pyspark_fake_cls, pyspark_fake_multi_cls,\n",
    "    pyspark_fake_seq_tag)\n",
    "from m3tl.special_tokens import EVAL, PREDICT, TRAIN\n",
    "from m3tl.utils import (get_embedding_table_from_model,\n",
    "                        set_is_pyspark, set_phase)\n",
    "\n",
    "\n",
    "class TestBase():\n",
    "    def __init__(self):\n",
    "        self.setUp()\n",
    "\n",
    "    def setUp(self) -> None:\n",
    "        self.tmpfiledir = tempfile.mkdtemp()\n",
    "        self.tmpckptdir = tempfile.mkdtemp()\n",
    "        self.prepare_params()\n",
    "\n",
    "    def tearDown(self) -> None:\n",
    "        shutil.rmtree(self.tmpfiledir)\n",
    "        shutil.rmtree(self.tmpckptdir)\n",
    "\n",
    "    def prepare_params(self):\n",
    "\n",
    "        self.problem_type_dict = {\n",
    "            'weibo_fake_ner': 'seq_tag',\n",
    "            # 'weibo_cws': 'seq_tag',\n",
    "            'weibo_fake_multi_cls': 'multi_cls',\n",
    "            'weibo_fake_cls': 'cls',\n",
    "            'weibo_masklm': 'masklm',\n",
    "            'weibo_fake_regression': 'regression',\n",
    "            'weibo_fake_vector_fit': 'vector_fit',\n",
    "            'weibo_premask_mlm': 'premask_mlm',\n",
    "            'fake_contrastive_learning': 'contrastive_learning'\n",
    "        }\n",
    "\n",
    "        self.processing_fn_dict = {\n",
    "            'weibo_fake_ner': get_weibo_fake_ner_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            # 'weibo_cws': get_weibo_cws_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            'weibo_fake_cls': get_weibo_fake_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            'weibo_fake_multi_cls': get_weibo_fake_multi_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            'weibo_masklm': get_weibo_masklm(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            'weibo_fake_regression': get_weibo_fake_regression(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            'weibo_fake_vector_fit': get_weibo_fake_vector_fit(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            'weibo_premask_mlm': get_weibo_premask_mlm(file_path='/data/m3tl/data/ner/weiboNER*'),\n",
    "            'fake_contrastive_learning': get_fake_contrastive_learning_fn(file_path='/data/m3tl/data/ner/weiboNER*')\n",
    "        }\n",
    "        self.params = Params()\n",
    "        self.params.tmp_file_dir = self.tmpfiledir\n",
    "        self.params.ckpt_dir = self.tmpckptdir\n",
    "        self.params.transformer_model_name = 'voidful/albert_chinese_tiny'\n",
    "        self.params.transformer_config_name = 'voidful/albert_chinese_tiny'\n",
    "        self.params.transformer_tokenizer_name = 'voidful/albert_chinese_tiny'\n",
    "        self.params.transformer_model_loading = 'TFAlbertModel'\n",
    "        self.params.transformer_tokenizer_loading = 'BertTokenizer'\n",
    "        self.params.transformer_config_loading = 'AlbertConfig'\n",
    "        self.params.print_data = False\n",
    "        self.params.enable_modal_type = True\n",
    "        self.params.num_cpus = 1\n",
    "\n",
    "        self.params.register_multiple_problems(\n",
    "            problem_type_dict=self.problem_type_dict, processing_fn_dict=self.processing_fn_dict)\n",
    "        self.params.assign_problem(\n",
    "            'weibo_fake_ner&weibo_fake_cls&weibo_fake_regression&weibo_fake_vector_fit|weibo_fake_multi_cls|weibo_masklm|weibo_premask_mlm|fake_contrastive_learning', base_dir=self.tmpckptdir)\n",
    "        self.params.assign_embedding_layer('duplicate_data_augmentation_embedding')\n",
    "        self.params.duplicate_data_aug_problems = ['fake_contrastive_learning']\n",
    "        set_phase(TRAIN)\n",
    "\n",
    "    def build_body(self):\n",
    "        if not hasattr(self, 'params'):\n",
    "            self.prepare_params()\n",
    "        self.test_body = BertMultiTaskBody(params=self.params)\n",
    "\n",
    "    def build_top(self):\n",
    "        if not hasattr(self, 'params'):\n",
    "            self.prepare_params()\n",
    "        if not hasattr(self, 'test_body'):\n",
    "            self.build_body()\n",
    "        input_embeddings = get_embedding_table_from_model(\n",
    "            self.test_body.bert.bert_model)\n",
    "        self.test_top = BertMultiTaskTop(\n",
    "            params=self.params, input_embeddings=input_embeddings)\n",
    "\n",
    "    def build_all_model(self):\n",
    "        if not hasattr(self, 'params'):\n",
    "            self.prepare_params()\n",
    "        self.all_model = BertMultiTask(params=self.params)\n",
    "\n",
    "    def get_dataset(self, mode=TRAIN):\n",
    "        if not hasattr(self, 'params'):\n",
    "            self.prepare_params()\n",
    "        dataset = train_eval_input_fn(\n",
    "            params=self.params, mode=TRAIN)\n",
    "        return dataset\n",
    "\n",
    "    def get_one_batch_input(self, mode=TRAIN):\n",
    "        dataset = self.get_dataset(mode=mode)\n",
    "        dataset = dataset.repeat()\n",
    "\n",
    "        one_batch_data = next(dataset.as_numpy_iterator())\n",
    "        return one_batch_data\n",
    "\n",
    "    def get_one_batch_body_model_output(self, mode=TRAIN):\n",
    "        one_batch_data = self.get_one_batch_input(mode=mode)\n",
    "        if not hasattr(self, 'test_body'):\n",
    "            self.build_body()\n",
    "        features, hidden_features = self.test_body(one_batch_data)\n",
    "        return features, hidden_features\n",
    "\n",
    "    def test_mtl_model(self, mtl_model: tf.keras.Model, include_top: bool, mode: str):\n",
    "\n",
    "        one_batch_body_output = self.get_one_batch_body_model_output(mode=mode)\n",
    "        if not include_top:\n",
    "            if not hasattr(self, 'test_top'):\n",
    "                self.build_top()\n",
    "            model_outputs = mtl_model(one_batch_body_output)\n",
    "            return self.test_top(model_outputs)\n",
    "        else:\n",
    "            return mtl_model(one_batch_body_output)\n",
    "\n",
    "    def test_loss_combination_strategy(self, loss_combination_strategy_name: str, run_fit=True):\n",
    "        if not hasattr(self, 'params'):\n",
    "            self.prepare_params()\n",
    "        self.params.assign_loss_combination_strategy(\n",
    "            loss_combination_strategy_name)\n",
    "\n",
    "        train_dataset = self.get_dataset(mode=TRAIN)\n",
    "        eval_dataset = self.get_dataset(mode=EVAL)\n",
    "        self.build_all_model()\n",
    "\n",
    "        self.all_model.compile()\n",
    "        if run_fit:\n",
    "            self.all_model.fit(\n",
    "                train_dataset,\n",
    "                validation_data=eval_dataset,\n",
    "                steps_per_epoch=1,\n",
    "                validation_steps=1,\n",
    "                epochs=2)\n",
    "\n",
    "    def create_fake_embedding_layer(self) -> tf.keras.layers.Embedding:\n",
    "        layer = tf.keras.layers.Embedding(input_dim=25000, output_dim=64)\n",
    "        layer.build((1, 1))\n",
    "        return layer\n",
    "\n",
    "    def test_embedding_layer(self, embedding_layer: tf.keras.Model, mode=TRAIN):\n",
    "        one_batch_data = self.get_one_batch_input(mode=mode)\n",
    "        return embedding_layer(one_batch_data, mode == TRAIN)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class PysparkTestBase(TestBase):\n",
    "    def __init__(self):\n",
    "        self.setUp()\n",
    "\n",
    "    def setUp(self) -> None:\n",
    "        set_is_pyspark(True)\n",
    "        self.tmpfiledir = tempfile.mkdtemp()\n",
    "        self.tmpckptdir = tempfile.mkdtemp()\n",
    "        self.tmppysparkdir = tempfile.mkdtemp()\n",
    "        self.prepare_params()\n",
    "\n",
    "    def prepare_params(self):\n",
    "\n",
    "        self.problem_type_dict = {\n",
    "            'pyspark_fake_seq_tag': 'seq_tag',\n",
    "            'pyspark_fake_multi_cls': 'multi_cls',\n",
    "            'pyspark_fake_cls': 'cls',\n",
    "        }\n",
    "\n",
    "        self.processing_fn_dict = {\n",
    "            'pyspark_fake_seq_tag': pyspark_fake_seq_tag,\n",
    "            'pyspark_fake_multi_cls': pyspark_fake_multi_cls,\n",
    "            'pyspark_fake_cls': pyspark_fake_cls,\n",
    "        }\n",
    "        self.params = Params()\n",
    "        self.params.tmp_file_dir = self.tmpfiledir\n",
    "        self.params.ckpt_dir = self.tmpckptdir\n",
    "        self.params.transformer_model_name = 'voidful/albert_chinese_tiny'\n",
    "        self.params.transformer_config_name = 'voidful/albert_chinese_tiny'\n",
    "        self.params.transformer_tokenizer_name = 'voidful/albert_chinese_tiny'\n",
    "        self.params.transformer_model_loading = 'TFAlbertModel'\n",
    "        self.params.transformer_tokenizer_loading = 'BertTokenizer'\n",
    "        self.params.transformer_config_loading = 'AlbertConfig'\n",
    "        self.params.print_data = False\n",
    "        self.params.enable_modal_type = True\n",
    "        self.params.num_cpus = 1\n",
    "        self.params.pyspark_output_path = self.tmppysparkdir\n",
    "\n",
    "        self.params.register_multiple_problems(\n",
    "            problem_type_dict=self.problem_type_dict, processing_fn_dict=self.processing_fn_dict)\n",
    "        self.params.assign_problem(\n",
    "            'pyspark_fake_seq_tag&pyspark_fake_multi_cls&pyspark_fake_cls', base_dir=self.tmpckptdir)\n",
    "        set_phase(TRAIN)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# prepare top dummy inputs\n",
    "# non-empty inputs\n",
    "def create_dummy_features_hidden_features(batch_size=1, hidden_dim=768, sample_features: dict = None, problem: str = None):\n",
    "    seq_len = sample_features['text_input_ids'].shape[1]\n",
    "    dummy_features = {}\n",
    "    for feature_name, feature in sample_features.items():\n",
    "        if len(feature.shape) > 1:\n",
    "            dummy_features[feature_name] = np.zeros(\n",
    "                shape=(batch_size, *feature.shape[1:]), dtype=feature.dtype)\n",
    "        else:\n",
    "            dummy_features[feature_name] = np.zeros(\n",
    "                shape=(batch_size), dtype=feature.dtype)\n",
    "    dummy_features['model_input_mask'] = np.ones_like(\n",
    "        dummy_features['text_input_ids'], dtype=np.int32)\n",
    "\n",
    "    dummy_features['{}_masked_lm_positions'.format(problem)] = np.zeros(\n",
    "        shape=(batch_size, 20), dtype=np.int64)\n",
    "    dummy_features['{}_masked_lm_ids'.format(problem)] = np.zeros(\n",
    "        shape=(batch_size, 20), dtype=np.int64)\n",
    "    dummy_features['{}_masked_lm_weights'.format(problem)] = np.ones(\n",
    "        shape=(batch_size, 20), dtype=np.int64)\n",
    "\n",
    "    dummy_hidden = {\n",
    "        'seq': np.zeros(shape=(batch_size, seq_len, hidden_dim), dtype=np.float32),\n",
    "        'pooled': np.zeros(shape=(batch_size, hidden_dim), dtype=np.float32),\n",
    "    }\n",
    "    return dummy_features, dummy_hidden\n",
    "\n",
    "\n",
    "def test_top_layer(top_class, problem: str, params: Params, sample_features: dict, hidden_dim: int, test_batch_size_list : list=None, **kwargs):\n",
    "    logger.debug('Testing {}'.format(top_class.__name__))\n",
    "    top_layer = top_class(params, problem, **kwargs)\n",
    "    if test_batch_size_list is None:\n",
    "        test_batch_size_list = [0,1,2]\n",
    "        \n",
    "    for batch_size in test_batch_size_list:\n",
    "        logger.debug('testing batch size {}'.format(batch_size))\n",
    "        # non empty test\n",
    "        d_f, d_h = create_dummy_features_hidden_features(\n",
    "            batch_size=batch_size, hidden_dim=hidden_dim, sample_features=sample_features, problem=problem)\n",
    "        set_phase(TRAIN)\n",
    "        _ = top_layer((d_f, d_h))\n",
    "        set_phase(EVAL)\n",
    "        _ = top_layer((d_f, d_h))\n",
    "        set_phase(PREDICT)\n",
    "        _ = top_layer((d_f, d_h))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.10 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
